import sys
sys.path.append('C:/Users/Hu/Dropbox/Research/PythonWork/Cancer/src/STAT/')

import numpy as np
import csv
from datetime import datetime
from time import strftime
import matplotlib.pyplot as plt
import math
import random
from scipy import stats
import ols


# When the appropriate changepoint is not known: 

# This is an python function to pick the optimal changepoint for a 
# continuous piecewise linear regression, based on a Sum of Squares Error (SSE) criterion:


def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def build_data_list(inputCSV):
    sKey = []
    fn = inputCSV
    f = open(inputCSV)
    #ra = csv.DictReader(file(fn), dialect="excel")
    ra = csv.DictReader(f, dialect="excel")
    
    for record in ra:
        #print record[ra.fieldnames[0]], type(record[ra.fieldnames[-1]])
        for item in ra.fieldnames:
            temp = float(record[item])
            sKey.append(temp)
    sKey = np.array(sKey)
    sKey.shape=(-1,len(ra.fieldnames))
    return sKey


def pickBreakpoint(response, predictor):
    bpChoices = np.sort(predictor)
    results = np.zeros((len(predictor)-1, 2))
    
    for i in range(len(predictor)-1):
        x2star = (predictor - bpChoices[i]) * np.greater(predictor, bpChoices[i])   
        tempPredictor = np.array(zip(predictor, x2star))
        tempmodel = ols.ols(response, tempPredictor,'y',['x1', 'x2'])
        results[i,0] = i
        results[i,1] = tempmodel.sse

    optBP = int(results[np.argmin(results, axis = 0)[1],0])
    print 'Optimal changepoint: ', bpChoices[optBP], ' with SSE = ', results[optBP, 1]

    x2star = (predictor - bpChoices[optBP]) * np.greater(predictor, bpChoices[optBP])
    optPredictor = np.array(zip(predictor, x2star))
    optmodel = ols.ols(response, optPredictor,'y',['x1', 'x2'])
    
    return bpChoices[optBP], results, optmodel, optmodel.b[0]+optmodel.b[1]*predictor+optmodel.b[2]*x2star


#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
    print '===================================================='
    print "begin at " + getCurTime()
    filePath = 'C:/_DATA/migration_census_2000/'
    file = filePath + 'segmented_data.csv'
    data = build_data_list(file)

    #tempmodel = ols.ols(data[:,0], data[:,1],'y',['x1'])
    optBP, results, optmodel, y_hat = pickBreakpoint(data[:,0], data[:,1])
    #print optmodel.summary(), optmodel.y
    plt.scatter(data[:,1], data[:,0], label = 'Original data', color ='b')
    plt.plot(data[:,1], y_hat, label = 'Fitted line', color ='g', linewidth=2.5)
    plt.axvline(x=optBP, linestyle ='--', color ='r', label='Breakpoint line', linewidth=1.5)
    plt.legend(loc=3)
    plt.show()
    
    #x2star = (data[:,1] - optBP) * np.greater(data[:,1], optBP)
    #optPredictor = np.array(zip(data[:,1], x2star))
    #optmodel = ols.ols(data[:,0], tempPredictor,'y',['x1', 'x2'])
    #print tempmodel.summary()
